"""
工具函数包
提供各种辅助功能
"""

import functools
import logging
import os
import uuid

from langfuse.decorators import langfuse_context

from src.utils.llm_health_checker import LLMHealthChecker


def setup_custom_logger(
    name: str, level_str: str = "INFO", is_dev: bool = False
) -> logging.Logger:
    """
    设置自定义日志记录器

    Args:
        name: 日志记录器名称
        level_str: 日志级别字符串
        is_dev: 是否为开发环境

    Returns:
        logging.Logger: 配置好的日志记录器
    """
    level = getattr(logging, level_str)
    logger = logging.getLogger(name)
    logger.setLevel(level)

    if not logger.handlers:
        # 创建一个处理程序，将日志输出到控制台
        handler = logging.StreamHandler()
        handler.setLevel(level)

        # 创建一个格式化器，定义日志的格式
        if is_dev:
            formatter = logging.Formatter(
                "%(levelname)s:%(name)s:%(lineno)d] %(message)s"
            )
        else:
            formatter = logging.Formatter(
                "I%(asctime)s %(process)d %(name)s:%(lineno)d] %(message)s",
                "%m%d %H:%M:%S.",
            )
        handler.setFormatter(formatter)

        # 将处理程序添加到日志记录器
        logger.addHandler(handler)

    return logger


def remove_trailing_slash(endpoint: str) -> str:
    """
    移除URL末尾的斜杠

    Args:
        endpoint: 输入URL

    Returns:
        str: 处理后的URL
    """
    return endpoint.rstrip("/") if endpoint.endswith("/") else endpoint


def init_langfuse(settings):
    """
    初始化Langfuse跟踪

    Args:
        settings: 应用设置
    """
    if settings.langfuse_enable:
        from langfuse import Langfuse

        logging.getLogger("wren-ai-service").info(
            f"LANGFUSE_ENABLE: {settings.langfuse_enable}"
        )
        logging.getLogger("wren-ai-service").info(
            f"LANGFUSE_HOST: {settings.langfuse_host}"
        )

        # 这些是可选的，如果没有设置，就不会初始化Langfuse客户端
        public_key = os.environ.get("LANGFUSE_PUBLIC_KEY", None)
        secret_key = os.environ.get("LANGFUSE_SECRET_KEY", None)

        if public_key and secret_key:
            Langfuse(
                public_key=public_key,
                secret_key=secret_key,
                host=settings.langfuse_host,
            )

            # 初始化langfuse上下文
            langfuse_context.configure(
                enabled=settings.langfuse_enable,
                host=settings.langfuse_host,
                public_key=public_key,
                secret_key=secret_key,
            )
        else:
            logging.getLogger("wren-ai-service").warning(
                "Langfuse keys not found, tracking is disabled"
            )
    else:
        logging.getLogger("wren-ai-service").info("Langfuse tracking is disabled")


def trace_metadata(func):
    """
    这个装饰器用于向当前的Langfuse跟踪添加元数据。
    它应该在创建跟踪后应用。以下是使用示例：

    ```python
    @observe(name="Mock")
    @trace_metadata
    async def mock():
        return "Mock"
    ```

    Args:
        func (Callable): 要装饰的函数

    Returns:
        Callable: 装饰后的函数
    """

    def extract(*args) -> dict:
        request = args[1]  # 修复请求对象的位置
        metadata = {}

        if hasattr(request, "project_id"):
            metadata["project_id"] = request.project_id
        if hasattr(request, "thread_id"):
            metadata["thread_id"] = request.thread_id
        if hasattr(request, "mdl_hash"):
            metadata["mdl_hash"] = request.mdl_hash
        if hasattr(request, "user_id"):
            metadata["user_id"] = request.user_id
        if hasattr(request, "query"):
            metadata["query"] = request.query

        return metadata

    @functools.wraps(func)
    async def wrapper(*args, **kwargs):
        trace_id = langfuse_context.get_current_trace_id()

        results = await func(*args, **kwargs, trace_id=trace_id)

        addition = {}
        if isinstance(results, dict):
            additional_metadata = results.get("metadata", {})
            addition.update(additional_metadata)

        metadata = extract(*args)
        service_metadata = kwargs.get(
            "service_metadata",
            {
                "pipes_metadata": {},
                "service_version": "",
            },
        )
        langfuse_metadata = {
            **service_metadata.get("pipes_metadata"),
            **addition,
            "mdl_hash": metadata.get("mdl_hash"),
            "project_id": metadata.get("project_id"),
            "query": metadata.get("query"),
        }
        langfuse_context.update_current_trace(
            user_id=metadata.get("user_id"),
            session_id=metadata.get("thread_id"),
            release=service_metadata.get("service_version"),
            metadata=langfuse_metadata,
        )

        return results

    return wrapper


def get_trace_id():
    """生成唯一的跟踪ID。"""
    return str(uuid.uuid4())


__all__ = [
    "setup_custom_logger",
    "init_langfuse",
    "LLMHealthChecker",
    "trace_metadata",
    "get_trace_id",
    "remove_trailing_slash",
]
